import uuid
import datetime
import re
import requests
from typing import List, Optional
from pymilvus import Collection, connections, model
from langgraph.graph import StateGraph
from typing import TypedDict
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

app = FastAPI()

MILVUS_HOST = "192.168.3.111"
MILVUS_PORT = "19530"
COLLECTION_NAME = "memory_record"
LLM_API_URL = "http://192.168.3.111:11434/api/chat"
EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"

connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
collection = Collection(COLLECTION_NAME)

embedding_fn = model.dense.SentenceTransformerEmbeddingFunction(
    model_name=EMBEDDING_MODEL_NAME, device="cuda:0"
)


class InputTextRequest(BaseModel):
    user_id: str
    input_text: str


class OptimizeState(TypedDict, total=False):
    user_id: str
    input_text: str
    optimized_text: Optional[str] = None
    event_time: Optional[str] = None
    category: Optional[str] = None


def optimize_text(input_text: str) -> (str, str, str):
    prompt = (
        "请将我提供的一段文字进行优化，使表达更通顺、简洁自然。同时，提取出描述的事件发生的时间（如果有明确时间的话；如果没有时间，请返回空字符串）。"
        "另外，请根据文本的内容给出一个分类标签（从这些标签中选择一个：童年与家乡, 家庭背景, 教育经历, 职业生涯, 婚姻与家庭, 子女教育, 人生挑战, 朋友与社交, 兴趣与爱好, 人生哲学, 社会变迁, 人生感悟）。"
        "最后请按照下面的格式返回：\n"
        "优化后的文本：在这里输出优化后的文字\n"
        "事件发生时间：在这里输出时间，如“2023年5月”或为空\n"
        "分类标签：在这里输出分类标签\n"
        f"[文字]\n{input_text}"
    )

    try:
        response = requests.post(
            LLM_API_URL,
            json={
                "model": "deepseek-r1:70b",
                "messages": [{"role": "user", "content": prompt}],
                "option": {"temperature": 0},
                "stream": False,
            }
        )
        response.raise_for_status()
        content = response.json().get("message", {}).get("content", "未能生成答案")
        content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL)

        print("模型返回的内容：", content)

        parts = content.split("优化后的文本：")
        if len(parts) < 2:
            raise ValueError(f"返回的内容格式不正确，无法从内容中提取出'优化后的文本'部分。返回内容：{content}")

        optimized_text = parts[1].split("事件发生时间：")[0].strip()

        event_time_section = parts[1].split("事件发生时间：")
        if len(event_time_section) < 2:
            event_time = ""
        else:
            event_time = event_time_section[1].split("分类标签：")[0].strip()

        category_section = event_time_section[1].split("分类标签：")
        if len(category_section) < 2:
            category = ""
        else:
            category = category_section[1].strip()

        return optimized_text, event_time, category
    except requests.exceptions.RequestException as e:
        print(f"请求错误: {e}")
        return "", "", ""
    except ValueError as ve:
        print(f"解析错误: {ve}")
        return "", "", ""


def insert_into_milvus(
        collection: Collection,
        original_text: str,
        optimized_text: str,
        embedding: List[float],
        user_id: str,
        relation: str = "",
        happen_time: Optional[str] = None,
        title: str = ""
):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    happen_time = happen_time or timestamp
    data_to_insert = [
        [str(uuid.uuid4())],
        [user_id],
        [original_text],
        [optimized_text],
        [relation],
        [timestamp],
        [happen_time],
        [title],
        [embedding],
    ]
    print("数据准备插入 Milvus：", data_to_insert)
    collection.insert(data_to_insert)


workflow = StateGraph(state_schema=OptimizeState)


def optimize_step(state: OptimizeState):
    optimized_text, event_time, category = optimize_text(state["input_text"])
    return {
        "optimized_text": optimized_text,
        "event_time": event_time,
        "category": category,
    }


workflow.add_node("optimize_step", optimize_step)


def save_step(state: OptimizeState):
    embedding = embedding_fn(state["input_content"])
    insert_into_milvus(
        collection=collection,
        original_text=state["input_text"],
        optimized_text=state["optimized_text"],
        embedding=embedding,
        user_id=state["user_id"],
        title=state["category"],
        happen_time=state["event_time"]
    )
    return {}




workflow.add_node("save_step", save_step)

workflow.set_entry_point("optimize_step")
workflow.add_edge("optimize_step", "save_step")
compiled_workflow = workflow.compile()


@app.post("/optimize_text/")
async def optimize_text_endpoint(request: InputTextRequest):
    initial_state = OptimizeState(user_id=request.user_id, input_text=request.input_text)

    try:
        final_state = compiled_workflow.invoke(initial_state)
        return {
            "optimized_text": final_state.get("optimized_text"),
            "event_time": final_state.get("event_time"),
            "category": final_state.get("category"),
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during optimization: {str(e)}")

if __name__ == "__main__":
    uvicorn.run("4:app", host="127.0.0.1", port=8001, reload=True)

